class Solution
{
public:
    // 向左、向右，记录小于k的数与大于k的数的个数的差值的出现次数，再匹配
    int countSubarrays(vector<int> &nums, int k)
    {
        int pos = find(nums.begin(), nums.end(), k) - nums.begin();
        int n = nums.size();
        int maxExtend = 2 * min(k - 1, n - k) + 1;
        int leftMaxExtend = min(maxExtend, pos);
        int rightMaxExtend = min(maxExtend, n - 1 - pos);
        unordered_map<int, int> leftDiffCount;
        int leftLowerCount = 0;
        for (int i = 0; i <= leftMaxExtend; ++i)
        {
            if (nums[pos - i] < k)
            {
                ++leftLowerCount;
            }
            ++leftDiffCount[leftLowerCount - (i - leftLowerCount)];
        }
        unordered_map<int, int> rightDiffCount;
        int rightLowerCount = 0;
        for (int j = 0; j <= rightMaxExtend; ++j)
        {
            if (nums[pos + j] < k)
            {
                ++rightLowerCount;
            }
            ++rightDiffCount[rightLowerCount - (j - rightLowerCount)];
        }

        int result = 0;
        for (auto &[diff, count] : leftDiffCount)
        {
            if (rightDiffCount.count(-diff))
            {
                result += count * rightDiffCount[-diff];
            }
            if (rightDiffCount.count(-diff - 1))
            {
                result += count * rightDiffCount[-diff - 1];
            }
        }
        return result;
    }

    //复杂度n^2，超时
    int countSubarraysDp(vector<int> &nums, int k)
    {
        int pos = find(nums.begin(), nums.end(), k) - nums.begin();
        int n = nums.size();
        int maxExtend = 2 * min(k - 1, n - k) + 1;
        int leftMaxExtend = min(maxExtend, pos);
        int rightMaxExtend = min(maxExtend, n - 1 - pos);

        // dp[i][j] 表示向左扩展i个, 向右扩展j个的子数组中，小于k的个数
        vector<vector<int>> dp(leftMaxExtend + 1, vector<int>(rightMaxExtend + 1));
        dp[0][0] = 0;
        int result = 1; // [0][0]情况对应1种

        auto checkSubArray = [&](int i, int j) -> void
        {
            int arrayLen = i + j + 1;
            int lowerCount = dp[i][j];
            if (lowerCount == (arrayLen - 1) / 2)
            {
                ++result;
            }
        };

        for (int i = 0; i <= leftMaxExtend; ++i)
        {
            if (i > 0)
            {
                dp[i][0] = dp[i - 1][0] + (nums[pos - i] < k);
                checkSubArray(i, 0);
            }
            for (int j = 1; j <= rightMaxExtend; ++j)
            {
                dp[i][j] = dp[i][j - 1] + (nums[pos + j] < k);
                checkSubArray(i, j);
            }
        }

        return result;
    }
};